#ifndef AST_HPP
#define AST_HPP

#include <vector>
#include <string>
using namespace std;

class Token;
class Visitor;
class SymbolTable;
class TypeTable;

class Context {
public:
    static const int Load = 0;
    static const int Store = 1;
};

class Node {
private:
    Node* _parent;

public:
    Node() : _parent(nullptr)   {}
    virtual void accept(Visitor* v) = 0;

    void set_parent(Node* p)    { _parent = p; }
    Node* parent()              { return _parent; }
};

enum OpType {
    AST_OP_ADD,
    AST_OP_SUB,
    AST_OP_MUL,
    AST_OP_DIV,
    AST_OP_MOD,
    AST_OP_BIT_AND,
    AST_OP_BIT_OR,
    AST_OP_BIT_XOR,
    AST_OP_LOG_AND,
    AST_OP_LOG_OR,
    AST_OP_LEFT_SHIFT,
    AST_OP_RIGHT_SHIFT,
    AST_OP_CMP,
};

class BinaryOp : public Node {
protected:
    OpType _op_type;
    Node* _left;
    Node* _right;
    
public:
    BinaryOp(OpType op_type, Node* left, Node* right) : 
        _op_type(op_type), _left(left), _right(right) {
        if (left) left->set_parent(this);
        if (right) right->set_parent(this);
    }

    virtual void accept(Visitor* v);

    void set_left(Node* left);
    void set_right(Node* right);

    OpType op_type()    { return _op_type; }
    Node* left()        { return _left; }
    Node* right()       { return _right; }
};

class ListNode : public Node {
private:
    vector<Node*>* _node_list;

public:
    ListNode();
    
    void add(Node* n);
    
    virtual void accept(Visitor* v);

    vector<Node*>* node_list()  { return _node_list; }
    int size()                  { return _node_list->size(); }

    vector<Node*>::iterator begin() { return _node_list->begin(); }
    vector<Node*>::iterator end()   { return _node_list->end(); }

    Node* get(int index)        { return (*_node_list)[index]; }
    void remove(Node* n);
    void insert(Node* n, Node* t);
};

class CmpNode : public BinaryOp {
private:
    int _cmp_op;

public:
    const static int CMP_LT = 0;
    const static int CMP_LE = 1;
    const static int CMP_EQU = 2;
    const static int CMP_NE = 3;
    const static int CMP_GT = 4;
    const static int CMP_GE = 5;


    CmpNode(int cmp_op) : BinaryOp(AST_OP_CMP, NULL, NULL), _cmp_op(cmp_op) {
    }

    CmpNode(Node* left, Node* right) : _cmp_op(-1), BinaryOp(AST_OP_CMP, left, right) {
    }

    int cmp_op()        { return _cmp_op; }
    virtual void accept(Visitor* v);
};

class ConstInt : public Node {
public:
    int _value;

    ConstInt(int v) : _value(v) {}

    virtual void accept(Visitor* v);
};

class ConstArray : public Node {
public:
    ListNode* arr;

    ConstArray(ListNode* v) : arr(v) {}

    virtual void accept(Visitor* v);
};

class ConstString : public Node {
public:
    char*   _value;
    int     _length;

    ConstString(int length, const char* s);

    ~ConstString();

    virtual void accept(Visitor* v);
};

class LogicNotNode: public Node {
private:
    Node* _value;
public:
    LogicNotNode(Node* v) : _value(v) {
    }

    Node* value() { return _value; }
    virtual void accept(Visitor* v);
};

class VarNode : public Node {
private:
    string  _name;
    ListNode*   _poly_args;

public:
    VarNode(Token* name);
    VarNode(const char* name) : _name(name) {}
    
    virtual void accept(Visitor* v);
    string& name()        { return _name; }

    void set_poly_args(ListNode* n) { _poly_args = n; }
    ListNode* poly_args()                { return _poly_args; }
};

class CallNode : public Node {
private:
    Node*       _name;
    ListNode*   _args;

public:
    CallNode(Node* name, ListNode* args) : _name(name), _args(args) {}
    
    virtual void accept(Visitor* v);

    Node* callee()      { return _name; }
    ListNode* args()    { return _args; }

    void set_callee(Node* n)   { _name = n; }
};

class SubscrNode : public Node {
private:
    Node*       _name;
    Node*       _index;
    int         _is_generic_inst;

public:
    SubscrNode(Node* name, Node* i) : _name(name), _index(i) {}
    
    virtual void accept(Visitor* v);

    Node* name()      { return _name; }
    Node* index()     { return _index; }

    void set_generic_inst(int f)    { _is_generic_inst = f; }
    int is_generic_inst()           { return _is_generic_inst; }
};

class AssignNode : public Node {
private:
    Node*   _left;
    Node*   _right;
    
public:
    AssignNode(Node* left, Node* right):
        _left(left), _right(right){}
        
    virtual void accept(Visitor* v);

    Node* left()        { return _left; }
    Node* right()       { return _right; }
};

class PrintNode : public Node {
private:
    Node*    _value;
    
public:
    PrintNode(Node* value) : _value(value) {}
        
    virtual void accept(Visitor* v);
    Node* value() { return _value; }
};

class ReturnNode : public Node {
private:
    Node*    _value;

public:
    ReturnNode(Node* value) : _value(value) {}

    virtual void accept(Visitor* v);
    Node* value() { return _value; }
};

class IfNode : public Node {
private:
    Node*   _cond;
    Node*   _then;
    Node*   _else;

public:
    IfNode(Node* cond, Node* then_part, Node* else_part) : 
        _cond(cond), _then(then_part), _else(else_part) {}

    virtual void accept(Visitor* v);

    Node* cond()        { return _cond; }
    Node* then_part()   { return _then; }
    Node* else_part()   { return _else; }
};

class WhileNode : public Node {
private:
    Node*   _cond;
    Node*   _body;

public:
    WhileNode(Node* cond, Node* body) :
        _cond(cond), _body(body) {}

    virtual void accept(Visitor* v);

    Node* cond()        { return _cond; }
    Node* body()        { return _body; }
};

enum TypeEnum {
    TYPE_NODE_ARR,
    TYPE_NODE_CON,
};

class ArrowTypeNode : public Node {
private:
    Node* _src;
    Node* _dst;

public:
    ArrowTypeNode(Node* src, Node* dst) : 
        _src(src), _dst(dst) {}

    Node* src()     { return _src; }
    Node* dst()     { return _dst; }

    virtual void accept(Visitor* v);
};

class ArrayTypeNode : public Node {
private:
    Node* _elm_type;
    Node* _size;

public:
    ArrayTypeNode(Node* t, Node* s) : 
        _elm_type(t), _size(s) {}

    Node* elm_type()    { return _elm_type; }
    Node* size()        { return _size; }

    virtual void accept(Visitor* v);
};

class VarDefNode : public Node {
private:
    string _name;
    Node*  _type;
    Node*  _init_value;

public:
    VarDefNode(string name, Node* var_type, Node* init_value) :
        _name(name), _type(var_type), _init_value(init_value) {}

    virtual void accept(Visitor* v);

    const string& name()    { return _name; }
    Node* def_type()        { return _type; }
    Node* init_value()      { return _init_value; }
};

class ParamDefNode : public Node {
private:
    string      _name;
    Node*   _type;
    Node*       _default;

public:
    ParamDefNode(string name, Node* var_type, Node* init_value) :
        _name(name), _type(var_type), _default(init_value) {}

    virtual void accept(Visitor* v);

    const string& name()        { return _name; }
    Node* def_type()            { return _type; }
    Node* default_value()       { return _default; }
};

class FuncDefNode : public Node {
private:
    string      _name;
    Node*       _return_type;
    ListNode*   _args;
    Node*       _body;
    ListNode*   _poly_params;

    ListNode*   _generic_args;

    SymbolTable*    _sym_table;
    TypeTable*      _type_table;

public:
    FuncDefNode(Token* t);
    FuncDefNode(string& name, ListNode* args, Node* ret_tp, Node* body): 
        _name(name), _args(args), _return_type(ret_tp), _body(body) {}

    virtual void accept(Visitor* v);

    const string& name()            { return _name; }
    ListNode* args()                { return _args; }
    void set_args(ListNode* args)   { _args = args; }
    Node* body()                    { return _body; }
    void set_body(Node* body)       { _body = body; }
    Node* ret_type()                { return _return_type; }
    void set_ret_type(Node* t)      { _return_type = t; }

    void set_generic_args(ListNode* args)   { _generic_args= args; }
    ListNode* generic_args()                { return _generic_args; }
    bool has_generic_args()                 { return _generic_args != nullptr && _generic_args->size() > 0; }

    SymbolTable* sym_table()            { return _sym_table; }
    void set_sym_table(SymbolTable* st) { _sym_table = st; }

    TypeTable* type_table()             { return _type_table; }
    void set_type_table(TypeTable* tt)  { _type_table = tt; }
};

class LambdaDef : public Node {
private:
    Node*           _param;
    Node*           _body;

    SymbolTable*    _sym_table;

public:
    LambdaDef(Node* param, Node* body):
        _param(param), _body(body){}
    ~LambdaDef();

    Node* param() { return _param; }
    Node* body()  { return _body; }

    SymbolTable* sym_table()            { return _sym_table; }
    void set_sym_table(SymbolTable* st) { _sym_table = st; }

    virtual void accept(Visitor* visitor);
};

/*
 * Nodes for type defination.
 */
class TypeVarNode : public Node {
private:
    string  _name;

public:
    TypeVarNode(string name) : _name(name) {}

    virtual void accept(Visitor* v);
    string& name()        { return _name; }
};

class TypeVarDefNode : public Node {
private:
    string  _name;
    Node*   _value;

public:
    TypeVarDefNode(string& name, Node* value) : _name(name), _value(value) {}
    ~TypeVarDefNode();

    string& name()          { return _name; }
    Node* value()            { return _value; }

    virtual void accept(Visitor* visitor);
};

class TypeFuncDefNode : public Node {
private:
    string  _name;
    Node*   _body;
    ListNode*   _args;

public:
    TypeFuncDefNode(string& name) : _name(name) {}
    ~TypeFuncDefNode();

    string& name()              { return _name; }
    Node* body()                { return _body; }
    void set_body(Node* b)      { _body = b; }
    ListNode* args()            { return _args; }
    void set_args(ListNode* a)  { _args = a; }

    virtual void accept(Visitor* visitor);
};

class TypeCallNode : public Node {
private:
    Node*       _name;
    ListNode*   _args;

public:
    TypeCallNode(Node* name, ListNode* args) : _name(name), _args(args) {}

    virtual void accept(Visitor* v);

    Node* callee()      { return _name; }
    ListNode* args()    { return _args; }
};

class TypeConstInt : public Node {
public:
    int value;

    TypeConstInt(int v) : value(v) {}

    virtual void accept(Visitor* v);
};

class BuiltinConst : public Node {
public:
    int _builtin_id;

    static const int UNIT = 0;
    static const int TRUE = 1;
    static const int FALSE = 2;
    
    static BuiltinConst* UnitNode;
    static BuiltinConst* TrueNode;
    static BuiltinConst* FalseNode;
    
    BuiltinConst(int v) : _builtin_id(v) {}
    
    virtual void accept(Visitor* v);
};

#endif

